import Image
import numpy as np
from itertools import product
import cma
from math import sin, cos, asin, acos, ceil
#from multiprocessing import Pool

#p = Pool(10)

def tri(x):
	return asin(cos(x))

def car(x):
	return ceil(cos(x))

def arrondi(x):
	return int(round(x))

def product2(degre, rep):
	res=[]
	for row in product(range(degre+1), repeat=2):
		if sum(row) <= degre:
			res.append(row)
	return res

def createPoly(params, degre):
	return lambda x : sum([params[i] * x[0]**powers[0] * x[1]**powers[1] for i, powers in enumerate(product2(degre, 2))])

def createFunc(params):
	return lambda x : createPoly(params[:6], 2)(x) * cos(createPoly(params[6:12], 2)(x))

def ecartCarre(y1, y2):
	y1 = np.array(y1)
	y2 = np.array(y2)
	n = len(y1)
	ecart = y1-y2
	ecart = ecart**2
	ecart = ecart.sum()
	return ecart/n

#x = i / width
#y = i % width
NB_PIX_INIT = 50
def getXY(i, width, height, idCompression):
	nbPix = NB_PIX_INIT * 2**(idCompression-1)
	lenMeanX = ceil(float(width)/nbPix)
	lenMeanY = ceil(float(height)/nbPix)
	x = i / width
	x -= x%lenMeanX
	x += lenMeanX/2
	y = i % width
	y -= y%lenMeanY
	y += lenMeanY/2
	return int(x), int(y)


imageFileName = "texture03nonUnif.jpg"
imageFileName = "texture02lisse.jpg"
imageFileName = "texture03lisse.jpg"
image = Image.open(imageFileName)
dataImage = list(image.getdata())
width = image.size[0]
height = image.size[1]

NB_PIX_INIT = 20
nbCompressions = 1
nbPix = NB_PIX_INIT
while nbPix < width or nbPix < height:
	nbCompressions += 1
	nbPix *= 2
red = [{"nbVal":0,"data":[]} for x in xrange(nbCompressions)]
green = [{"nbVal":0,"data":[]} for x in xrange(nbCompressions)]
blue = [{"nbVal":0,"data":[]} for x in xrange(nbCompressions)]
dataReg = [[] for x in xrange(nbCompressions)]

for i, (r, g, b) in enumerate(dataImage):
	for j in xrange(nbCompressions):
		x,y = getXY(i, width, height, j+1)
		if [x,y] not in dataReg[j]:
			dataReg[j].append([x, y])
		idRow = dataReg[j].index([x,y])
		if idRow == len(red[j]["data"]):
			red[j]["data"].append(0)
			green[j]["data"].append(0)
			blue[j]["data"].append(0)
		red[j]["data"][idRow] += r
		green[j]["data"][idRow] += g
		blue[j]["data"][idRow] += b
		red[j]["nbVal"] += 1
		green[j]["nbVal"] += 1
		blue[j]["nbVal"] += 1
	#x = i / width
	#y = i % width
	#dataReg.append([x, y]) #, cos(x), cos(y)])
	#red.append(r)
	#green.append(g)
	#blue.append(b)

for j in xrange(nbCompressions):
	dataReg[j] = np.array(dataReg[j])
	for i in xrange(len(red[j])):
		red[j]["data"][i] /= red[j]["nbVal"]
		green[j]["data"][i] /= green[j]["nbVal"]
		blue[j]["data"][i] /= blue[j]["nbVal"]
	red[j]["data"] = np.array(red[j]["data"])
	green[j]["data"] = np.array(green[j]["data"])
	blue[j]["data"] = np.array(blue[j]["data"])

dimF = 12
idCompression = 0

def fObj(x):
	xRed = x[:dimF]
	xGreen = x[dimF:dimF*2]
	xBlue = x[dimF*2:dimF*3]
	fRed = createFunc(xRed)
	fGreen = createFunc(xGreen)
	fBlue = createFunc(xBlue)
	redEstime = map(fRed, dataReg[idCompression])
	greenEstime = map(fGreen, dataReg[idCompression])
	blueEstime = map(fBlue, dataReg[idCompression])
	ecartRed = ecartCarre(red[idCompression]["data"], redEstime)
	ecartGreen = ecartCarre(green[idCompression]["data"], greenEstime)
	ecartBlue = ecartCarre(blue[idCompression]["data"], blueEstime)
	ecart = ecartRed + ecartGreen + ecartBlue
	print "DEBUG : fObj evaluated", idCompression, ecart
	print "Sol :", x
	return ecart
	
xOpts = []
for j in xrange(nbCompressions):
	for xOpt in xOpts:
		xRed = xOpt[:dimF]
		xGreen = xOpt[dimF:dimF*2]
		xBlue = xOpt[dimF*2:dimF*3]
		fRed = createFunc(xRed)
		fGreen = createFunc(xGreen)
		fBlue = createFunc(xBlue)
		redEstime = map(fRed, dataReg[idCompression])
		greenEstime = map(fGreen, dataReg[idCompression])
		blueEstime = map(fBlue, dataReg[idCompression])
		red[idCompression]["data"] -= redEstime
		green[idCompression]["data"] -= greenEstime
		blue[idCompression]["data"] -= blueEstime
	params0 = [0 for x in xrange(dimF*3)]
	sd0 = 10
	res_cma = cma.fmin(fObj, params0, sd0)
	print "etape", idCompression, "finished"
	idCompression += 1
	xOpts.append(res_cma[0])
idCompression -= 1

redEstime = np.zeros((width*height,1))
greenEstime = np.zeros((width*height,1))
blueEstime = np.zeros((width*height,1))

for xOpt in xOpts:
	xRed = xOpt[:dimF]
	xGreen = xOpt[dimF:dimF*2]
	xBlue = xOpt[dimF*2:dimF*3]
	fRed = createFunc(xRed)
	fGreen = createFunc(xGreen)
	fBlue = createFunc(xBlue)
	redEstime += map(fRed, dataReg[idCompression])
	greenEstime += map(fGreen, dataReg[idCompression])
	blueEstime += map(fBlue, dataReg[idCompression])

red[idCompression]["data"] = map(arrondi, redEstime)
green[idCompression]["data"] = map(arrondi, greenEstime)
blue[idCompression]["data"] = map(arrondi, blueEstime)
dataNewImage = zip(redEstime, greenEstime, blueEstime)
newImage = Image.new(image.mode, image.size) 
newImage.putdata(dataNewImage) 

outImageFileName = imageFileName[:-4] + "_opt02" + imageFileName[-4:]
newImage.save(outImageFileName)
